import time
import json
import hmac
import base64
import websocket
from okx import Trade,MarketData
from datetime import datetime

api_key = '71db79ae-81f6-4e09-8541-cf210849a70b'
secret_key = '1BB9FBDDFCF31DC43012B2C7A88AFDD1'
# api_key = '6860449c-c557-434b-9eff-2c2546959cc7'
# secret_key = '4F532857977CA2410D335AB85ED846A5'
passphrase = 'Pfy@132465'
private_ws_url = "wss://ws.okx.com:8443/ws/v5/private"
trade_api = Trade.TradeAPI(api_key, secret_key, passphrase,flag='0')
market_api = MarketData.MarketAPI(flag='0')

# 交易对和网格参数
symbol = 'SOL-USDT'  # 交易对
lower_price = 142  # 网格最低价格
upper_price = 152  # 网格最高价格
amount_per_order = 0.05  # 每笔订单的数量
all_price = []
while lower_price <= upper_price:
    all_price.append(lower_price)
    lower_price += 0.5
print(all_price)
orderID = [str(i+100000000) for i in range(1,len(all_price)+1)]
print(orderID)

grid_order = {
    'instId': symbol,
    'tdMode': 'cash',
    'ordType': 'limit',
    'sz': str(amount_per_order),
}

# 通过最新价格来计算基准货币仓位
def getPosition(lastPrice):
    pos = 0
    for i in all_price:
        if lastPrice < i:
            pos += 1
    return pos * amount_per_order

# 创建网格订单
def create_grid_orders(lastPrice):
    # 1、建仓
    grid_order["side"] = "buy"
    grid_order["clOrdId"] = "100000000"
    grid_order["ordType"] = "limit"
    grid_order["sz"] = str(getPosition(lastPrice))
    grid_order["px"] = str(lastPrice+2)
    response = trade_api.place_order(**grid_order)
    print(f"Order executed: {response}")
    grid_order["ordType"] = "limit"
    grid_order["sz"] = str(amount_per_order)
    time.sleep(1)
    print("========================================")
    # 2、挂单
    for i in range(len(all_price)):
        if all_price[i] < lastPrice:
            grid_order["side"] = "buy"
            grid_order["clOrdId"] = orderID[i]
            grid_order["px"] = all_price[i]
            response = trade_api.place_order(**grid_order)
            print(f"Order executed: {response}")
        else:
            grid_order["side"] = "sell"
            grid_order["clOrdId"] = orderID[i]
            grid_order["px"] = all_price[i]
            response = trade_api.place_order(**grid_order)
            print(f"Order executed: {response}")


def generate_signature(api_key, secret_key, passphrase):
    timestamp = str(time.time())
    message = timestamp + 'GET' + '/users/self/verify'
    mac = hmac.new(bytes(secret_key, 'utf-8'), bytes(message, 'utf-8'), digestmod='sha256')
    sign = base64.b64encode(mac.digest()).decode('utf-8')
    return timestamp, sign

def on_error(ws, error):
    print(f"Error: {error}")

def on_close(ws, close_status_code, close_msg):
    print(f"WebSocket closed with code: {close_status_code}, message: {close_msg}")


def on_open(ws):
    print("WebSocket opened")
    # 生成登录签名
    timestamp, sign = generate_signature(api_key, secret_key, passphrase)

    # 发送登录请求
    login_message = {
        "op": "login",
        "args": [
            {
                "apiKey": api_key,
                "passphrase": passphrase,
                "timestamp": timestamp,
                "sign": sign
            }
        ]
    }
    ws.send(json.dumps(login_message))
    time.sleep(1)
    subscribe_message = {
        "op": "subscribe",
        "args": [
            {
                "channel": "orders",
                "instType": "SPOT",  # 现货交易
                "instId": symbol
            },
            {
                "channel": "account",
                "ccy": "USDT"  # 现货交易
            }
        ]
    }
    ws.send(json.dumps(subscribe_message))


def on_message(ws, message):
    data = json.loads(message)
    if ('arg' in data) and (data['arg']['channel'] == 'orders'):
        if ('data' in data) and (data['data'][0]['state'] == 'filled') and (data['data'][0]['clOrdId']!='100000000'):
            grid_order["clOrdId"] = data['data'][0]['clOrdId']
            grid_order["px"] = data['data'][0]['px']
            if data['data'][0]['side'] == 'buy':
                grid_order["side"] = "sell"
                response = trade_api.place_order(**grid_order)
                print(f"Order executed: {response}")
            if data['data'][0]['side'] == 'sell':
                grid_order["side"] = "buy"
                response = trade_api.place_order(**grid_order)
                print(f"Order executed: {response}")
        

def start_websocket():
    ws = websocket.WebSocketApp(private_ws_url,
                                on_message=on_message,
                                on_error=on_error,
                                on_close=on_close)
    ws.on_open = on_open
    ws.run_forever()

# 开始网格交易
if __name__ == '__main__':
    solPrice = market_api.get_ticker('SOL-USDT')['data'][0]['last']
    print(float(solPrice))
    create_grid_orders(float(solPrice))
    start_websocket()
    
